import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf

from sklearn.datasets import load_diabetes

data = load_diabetes()
print(data.keys())
x = data.data
y = data.target.reshape(-1, 1)

xy = np.c_[x, y]
import pandas as pd
import seaborn as sbn
df = pd.DataFrame(xy)
sbn.heatmap(df.corr(), annot=True)
plt.show()

x = x[:, 2:3]
print(x.shape, y.shape)

plt.scatter(x, y)
plt.show()

m, n = x.shape

train_size = 0.7
divide = int(m * train_size)
x_train, x_test = np.split(x, [divide])
y_train, y_test = np.split(y, [divide])

# hyper parameters
iters =20000
learning_rate = 0.1

class lrModel:

    def __init__(self, sess, name):
        self.sess = sess
        self.name = name
        self._build_net()

    def _build_net(self):
        # input place holders
        self.x = tf.placeholder(tf.float32, [None, n])
        self.y = tf.placeholder(tf.float32, [None, 1])

        W = tf.Variable(tf.random_normal([n, 1]))
        b = tf.Variable(tf.random_normal([1]))

        self.h = tf.matmul(self.x, W) + b

        # define cost/loss & optimizer
        self.cost = tf.reduce_mean(tf.square(tf.subtract(self.h, self.y)))
        self.optimizer = tf.train.AdamOptimizer(
            learning_rate=learning_rate).minimize(self.cost)

    def train(self, x_data, y_data):
        return self.sess.run([self.cost, self.optimizer],
                             feed_dict={self.x: x_data, self.y: y_data})

    def test(self, x_test):
        return self.sess.run(self.h, feed_dict={self.x: x_test})

# initialize
sess = tf.Session()
m1 = lrModel(sess, "m1")

sess.run(tf.global_variables_initializer())

print('Learning Started!')

# train my model
cost_his = []
for i in range(iters):
    cost_val, _ = m1.train(x_train, y_train)
    if i % 500 == 0:
        print("step", i, "cost:", cost_val)
        cost_his.append(cost_val)

plt.plot(cost_his[1:])
plt.show()

print('Learning Finished!')

# Test model
h = m1.test(x_test)
print('h_test:', h)

plt.scatter(y_test, h)
plt.plot(y_test, y_test, 'rx')
plt.show()

# plot fit curve
plt.scatter(x_test, y_test)
plt.plot(x_test, h, 'r-')
plt.show()

# compute R2
mse = np.mean(np.square(y_test - h))
ysigma = np.mean(np.square((y_test - np.mean(y_test))))
R2 = 1 - mse / ysigma
print("R2:", R2)
